# 作者: 王道 龙哥
# 2024年06月14日15时03分01秒
# dartou@qq.com
RED = 0 #类似于 C语言 的全局变量
BLACK = 1


class RedBlackNode:
    def __init__(self, value):
        self.value = value
        self.left = None
        self.right = None
        self.p = None
        self.color = RED


class RedBlackTree:
    def __init__(self):
        self.root = None

    def left_rotate(self, x: RedBlackNode):
        y: RedBlackNode = x.right
        x_parent: RedBlackNode = x.p
        if x_parent is None:  # x是原有的根，y称为新的根
            self.root = y
        elif x_parent.left is x:  # x在父亲左边
            x_parent.left = y
        else:  # x在父亲右边
            x_parent.right = y
        y.p = x_parent

        if y.left:  # 处理y的左子树
            y.left.p = x  # 左子树的父亲变为x
        x.right = y.left  # y的左子树变为x的右子树
        #
        y.left = x  # x变为y的左孩子
        x.p = y  # y变为x的父亲

    def right_rotate(self, node):
        if not node.left:
            return False
        node_left = node.left
        node_left.p = node.p
        if not node.p:
            self.root = node_left
        elif node == node.p.left:
            node.p.left = node_left
        elif node == node.p.right:
            node.p.right = node_left
        if node_left.right:
            node_left.right.p = node
        node.left = node_left.right
        node.p = node_left
        node_left.right = node

    def insert(self, value):
        node = RedBlackNode(value)
        if not self.root:
            self.root = node
        else:
            x = self.root
            while x:
                parent = x
                if node.value < x.value:
                    x = x.left
                else:
                    x = x.right
            # 判断node要放在parent左边还是右边
            if parent.value > node.value:
                parent.left = node
            else:
                parent.right = node
            node.p = parent
        self.insert_fixup(node)

    def insert_fixup(self, node):
        parent: RedBlackNode = node.p
        while parent and parent.color == RED:
            gprent: RedBlackNode = parent.p
            if gprent.left is parent:
                # 情形三
                uncle: RedBlackNode = gprent.right
                if uncle and uncle.color == RED:
                    gprent.color = RED
                    parent.color = BLACK
                    uncle.color = BLACK
                    node = gprent
                    parent = node.p
                    continue
                if parent.right is node:  # 情形四
                    self.left_rotate(parent)
                # 情形五
                self.right_rotate(gprent)
                gprent.color = RED
                parent.color = BLACK
            else:
                # 情形三
                uncle: RedBlackNode = gprent.left
                if uncle and uncle.color == RED:
                    gprent.color = RED
                    parent.color = BLACK
                    uncle.color = BLACK
                    node = gprent
                    parent = node.p
                    continue
                if parent.left is node:  # 情形四
                    self.right_rotate(parent)
                # 情形五
                self.left_rotate(gprent)
                gprent.color = RED
                parent.color = BLACK
        self.root.color = BLACK

    def mid_order(self, node: RedBlackNode):
        if node:
            self.mid_order(node.left)
            print(node.value, end=' ')
            self.mid_order(node.right)

    def rbtree_print(self, node, key, direction): #key是当前结点父亲的value
        if node:
            if direction == 0:  # tree是根节点
                print("%2d(B) is root" % node.value)
            else:  # tree是分支节点
                print("%2d(%s) is %2d's %6s child" % (
                    node.value, ("B" if node.color == 1 else "R"), key, ("right" if direction == 1 else "left")))

            self.rbtree_print(node.left, node.value, -1)
            self.rbtree_print(node.right, node.value, 1)


if __name__ == '__main__':
    number_list = (7, 4, 1, 8, 5, 2, 9, 6, 3)
    tree = RedBlackTree()
    for number in number_list:
        tree.insert(number)
    # 中序遍历进行打印
    # tree.mid_order(tree.root)
    tree.rbtree_print(tree.root,0,0)
